import sklearn
import os
import random
from PIL import Image
import torch
from loguru import logger
from transformers import AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
import time
SEVER = True
if SEVER:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu


class medi_info():
    def __init__(self):
        '''
        初始化：定义疾病的医学描述和可视化特征
        '''
        self.disease_descriptions = {
            "mild Diabetic Retinopathy": {
                "medical": "视网膜出现少量微小出血点和血管扩张，可能有微动脉瘤，视网膜可能出现轻微水肿。",
                "visual": "可能看到少量红色点状斑块，血管轻度弯曲或膨胀，图像中可能有细微的水肿或黄色小斑点。"
            },
            "moderate Diabetic Retinopathy": {
                "medical": "视网膜上有多个微小出血、渗出，血管壁可能硬化，部分区域可能缺血，可能伴随视网膜水肿。",
                "visual": "图像中可能看到明显的红色或黄色区域，血管可能不规则或扭曲，某些区域可能呈现轻微的膨胀或水肿表现。"
            },
            "proliferative Diabetic Retinopathy": {
                "medical": "新生血管增生，可能出血、脱落，造成视力丧失，可能伴有视网膜脱离和纤维血管膜形成。",
                "visual": "图像中可能有大量不规则血管线条，或者暗红色阴影区域，某些区域可能看到视网膜出现扭曲或水肿性改变。"
            },
            "early Glaucoma": {
                "medical": "眼压略高，视神经有轻微损伤，视神经盘边缘可能变浅，视网膜神经纤维层可能出现轻微萎缩。",
                "visual": "图像变化不明显，可能看到视网膜或视神经边缘轻微变浅，中心略模糊或边缘轻微变暗。"
            },
            "moderate Glaucoma": {
                "medical": "视野出现缺损，视神经萎缩，视网膜神经纤维层萎缩加重，部分区域可能出现视网膜神经盘的损伤。",
                "visual": "图像可能边缘模糊，出现视野凹陷或亮度不均，某些区域可能看到视神经盘的颜色或形态变化。"
            },
            "advanced Glaucoma": {
                "medical": "视神经严重损害，视力大幅下降，视神经可能出现显著萎缩，视网膜中央区域可能出现明显的损伤。",
                "visual": "图像中央可能出现暗区，整体模糊或失真严重，可能有视神经萎缩的明显标志，周边视网膜变化较为明显。"
            },
            "dry Macular Degeneration": {
                "medical": "黄斑部轻度退化，视力逐渐下降。",
                "visual": "中心可能出现颜色不均、轻微退化斑块。"
            },
            "wet Macular Degeneration": {
                "medical": "新生血管导致黄斑部出血、渗液，视力快速丧失。",
                "visual": "中心可能出现红斑、液体隆起或不规则明暗对比。"
            }
        }

    def get_description(self, disease_name):
        if disease_name in self.disease_descriptions:
            return self.disease_descriptions[disease_name]
        else:
            raise NotImplementedError(f"无法获取 {disease_name} 的病灶描述")


class janus_interface():
    def __init__(self, model_path, server=False):
        self.medi_info = medi_info()
        self.device = self.setup_environment(server)
        self.processor, self.model = self.load_model(model_path)
        print(f"接口初始化完成，开始工作")

    def setup_environment(self, server=True):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f'环境初始化完成，当前设备：{device}')
        return device

    def load_model(self, model_path):
        logger.info('开始加载 processor 和模型')
        processor = VLChatProcessor.from_pretrained(model_path)
        vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
            model_path, trust_remote_code=True
        )
        vl_gpt = vl_gpt.to(torch.float16).cuda().eval()
        logger.info('模型和 processor 加载完成')
        return processor, vl_gpt

    def predict(self, images, question):
        logger.info('开始模型预测')

        conversation = [
            {"role": "<|User|>", "content": f"<image_placeholder>\n{question}"},
            {"role": "<|Assistant|>", "content": ""}
        ]

        processor_inputs = self.processor(
            conversations=conversation, images=images, force_batchify=True
        ).to(self.model.device)

        processor_inputs = dict(processor_inputs)
        processor_inputs["pixel_values"] = processor_inputs["pixel_values"].to(
            dtype=torch.float16)

        with torch.cuda.amp.autocast():
            inputs_embeds = self.model.prepare_inputs_embeds(
                **processor_inputs)

            outputs = self.model.language_model.generate(
                inputs_embeds=inputs_embeds,
                attention_mask=processor_inputs["attention_mask"],
                pad_token_id=self.processor.tokenizer.eos_token_id,
                bos_token_id=self.processor.tokenizer.bos_token_id,
                eos_token_id=self.processor.tokenizer.eos_token_id,
                max_new_tokens=512,
                do_sample=False,
                use_cache=True,
            )

        response = self.processor.tokenizer.decode(
            outputs[0], skip_special_tokens=True)
        logger.info(f"模型预测结果: {response}")
        return response

    def crop_image(self, image_path, box_list):
        image = Image.open(image_path)
        cropped_images = [image.crop(box) for box in box_list]
        return cropped_images

    def preprocess_one(self, image_path, diagnosis_label, region_list=None):
        if region_list:
            cropped_images = self.crop_image(image_path, region_list)
        else:
            cropped_images = [Image.open(image_path)]

        desc = self.medi_info.get_description(diagnosis_label)
        medical_desc = desc["medical"]
        visual_desc = desc["visual"]

        question = (
            f"该患者的整张眼底图像的初步诊断为：{diagnosis_label}。"
            f"医学症状为：{medical_desc}"
            f"你可能在图像中看到：{visual_desc}"
            f"目前展示的是该图像的一小部分区域，请仔细观察该区域，判断上述症状是否出现在这一局部图像中。"
            f"请以如下格式作答：\n"
            f"<描述症状>（请简要描述你在图中看到的现象）</描述症状><结论>是 或 否</结论>"
        )
        question="请详细描述图片" # simple question for debug
        return cropped_images, question



def main():
    # 记录程序启动时间
    start_time = time.time()

    image_path = './test.png'
    diagnosis_label = "moderate Diabetic Retinopathy"
    region_list = [[110, 110, 170, 170], [30, 30, 90, 90], [10, 10, 70, 70], [160, 160, 220, 220]]

    # 记录模型加载时间
    model_load_start = time.time()
    processor = janus_interface(
        model_path='./pretrained/eyekowner',
        server=SEVER
    )
    model_load_end = time.time()
    model_load_time = model_load_end - model_load_start
    print(f"模型加载时间: {model_load_time:.2f} 秒")

    # 生成问题的时间
    preprocess_start = time.time()
    cropped_images, question = processor.preprocess_one(image_path, diagnosis_label, region_list)
    preprocess_end = time.time()
    preprocess_time = preprocess_end - preprocess_start
    print(f"预处理时间（生成问题和裁剪区域）: {preprocess_time:.2f} 秒")

    # 输出目录
    output_dir = './experiments/'
    os.makedirs(output_dir, exist_ok=True)

    # 遍历每个区域，并统计处理时间
    for idx, image in enumerate(cropped_images):
        region_start = time.time()

        # 预测
        answer = processor.predict([image], question)

        # 保存图像
        save_path = os.path.join(output_dir, f"{idx}.png")
        image.save(save_path)

        region_end = time.time()
        region_time = region_end - region_start

        # 输出时间统计
        print(f"Region {idx} 处理时间: {region_time:.2f} 秒")
        print(f"Image {idx} saved to: {save_path}")
        print(f"Answer: {answer}")

    # 计算整体运行时间
    total_end = time.time()
    total_time = total_end - start_time
    print(f"总运行时间: {total_time:.2f} 秒")

if __name__ == '__main__':
    main()
